import torch
import torch.nn as nn
import torch.optim as optim

from DatasetInit import *
from SimpleRNN   import *

from torch.utils.data import DataLoader


# 假设输入数据为一个 batch，大小为 (batch_size, seq_len, input_size)
seq_len = 1
num_samples = 15  # 生成 15 个样本
batch_size = 1

# 模型参数
input_size = 5  # 输入特征维度
hidden_size = 96  # 隐藏层维度
output_size = 10  # 输出大小

output_size_target = 16  # 输出大小

bitnum = 16

# 创建数据集实例
dataset = CustomDatasetTest(num_samples,seq_len,input_size,output_size_target)

# 创建 DataLoader
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

print("testing data ok")


# 创建模型
model = SimpleRNN(input_size, hidden_size, output_size)

#model = nn.DataParallel(model)
#model = model.cuda()



model = torch.load("model/full_model.pth")
print("Full model loaded.")

model.eval()

datasetsin = []
datasetsout = []

# 获取一个批次的数据
for batch_inputs, batch_outputs in dataloader:

    batch_inputs = torch.reshape(batch_inputs,[batch_size, seq_len, input_size])
    batch_outputs = torch.reshape(batch_outputs,[batch_size, seq_len, output_size_target])

    with torch.no_grad():

        test_input0 =  torch.reshape(batch_inputs,[1,input_size])

        target_output =  torch.reshape(batch_outputs,[1,output_size_target])
        print(f"Decimal: {test_input0.cpu().numpy()}, target Binary: {target_output.cpu().numpy()}")


        test_output = model(test_input0)
        pred_output = torch.abs(torch.ceil(test_output-0.5))

        for t in range(bitnum-1):

            test_output = torch.reshape(test_output,[2,input_size])
            test_input = torch.reshape(test_output[0],[1,input_size])


            test_output = model(test_input)

            #pred_output.append(torch.abs(torch.ceil(test_output-0.5)))

            pred_output = torch.cat((pred_output,torch.abs(torch.ceil(test_output-0.5))),dim=0)

        
        pred_output = torch.reshape(pred_output,[output_size_target,output_size])

        pred_output = pred_output.cpu().numpy()
        print(f"Decimal: {test_input0.cpu().numpy()}, pred Binary:\n {pred_output}")